#!/usr/bin/env python
"""
播放种子5201的音色测试
"""
import sys
import torch
from pathlib import Path

project_root = Path(__file__).parent
sys.path.insert(0, str(project_root))

print("=" * 70)
print(" " * 20 + "音色5201测试播放")
print("=" * 70)
print()

# 导入ChatTTS
print("正在加载ChatTTS...")
import ChatTTS

# 初始化
chat = ChatTTS.Chat()
chat.load(source='huggingface', compile=False, device='cpu')
print("✓ ChatTTS加载完成")
print()

# 使用种子5201生成音色
print("正在生成音色 (种子: 5201)...")
torch.manual_seed(5201)
spk_emb = chat.sample_random_speaker()
print("✓ 音色生成完成")
print()

# 测试文本
test_text = "大家好，我是小爱，一个活泼可爱的AI虚拟主播。很高兴认识大家！这是种子5201的音色测试。"

print(f"测试文本: {test_text}")
print()
print("正在合成语音...")

# 合成音频
params_infer_code = ChatTTS.Chat.InferCodeParams(
    spk_emb=spk_emb,
    temperature=0.3,
    top_P=0.7,
    top_K=20,
)

wavs = chat.infer([test_text], params_infer_code=params_infer_code)

if wavs is not None and len(wavs) > 0:
    print("✓ 语音合成完成")
    print()
    
    # 保存WAV文件
    import scipy.io.wavfile as wavfile
    output_file = project_root / "voice_5201_test.wav"
    
    # ChatTTS返回的是numpy数组
    import numpy as np
    audio_data = np.array(wavs[0])
    
    # 确保数据格式正确
    if audio_data.dtype != np.int16:
        # 转换为16位整数格式
        audio_data = (audio_data * 32767).astype(np.int16)
    
    # 保存为WAV文件 (采样率24000Hz)
    wavfile.write(str(output_file), 24000, audio_data)
    
    print(f"✓ 音频已保存: {output_file}")
    print(f"  采样率: 24000 Hz")
    print()
    
    # 自动播放
    print("正在播放音频...")
    try:
        import pygame
        pygame.mixer.init()
        pygame.mixer.music.load(str(output_file))
        pygame.mixer.music.play()
        
        # 等待播放完成
        import time
        while pygame.mixer.music.get_busy():
            time.sleep(0.1)
        
        print("✓ 播放完成")
    except Exception as e:
        print(f"⚠️ 自动播放失败: {e}")
        print(f"请手动播放文件: {output_file}")
    
    print()
    print("=" * 70)
    print("如果这个音色正确，配置文件已经设置为使用此音色。")
    print("如果不正确，我可以帮您生成其他种子的音色。")
    print("=" * 70)
else:
    print("✗ 语音合成失败")
